# Copyright (c) 2024 Huawei Technologies Co., Ltd.
#
# openMind is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

from typing import List, Optional

from transformers import DataCollatorForLanguageModeling, Trainer, TrainerCallback

from openmind.utils import get_logger
from openmind.flow.model import get_model, get_tokenizer
from openmind.flow.datasets import get_template, get_dataset_module
from openmind.flow.arguments import get_args


logger = get_logger(__name__)


def run_pt(
    callbacks: Optional[List["TrainerCallback"]] = None,
):
    tokenizer = get_tokenizer()
    model = get_model()
    template = get_template()
    dataset_module = get_dataset_module(tokenizer, template)

    args = get_args()
    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    trainer = Trainer(
        model=model,
        args=args.hf_seq2seq_args,
        tokenizer=tokenizer,
        data_collator=data_collator,
        callbacks=callbacks,
        **dataset_module,
    )

    if args.do_train:
        logger.info_rank0("Start training.")
        train_result = trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
        trainer.save_model()
        trainer.log_metrics("train", train_result.metrics)
        trainer.save_metrics("train", train_result.metrics)
        trainer.save_state()
